import time
from concurrent.futures import Future
from concurrent.futures import ThreadPoolExecutor
from selectors import DefaultSelector, EVENT_READ, EVENT_WRITE
from typing import Callable, Any
import socket

def thread_try(fn):
    def in_fn(*args, **kwargs):
        try:
            return fn(*args, **kwargs)
        except Exception as e:
                print(f'thread_try err {e}')
                return None
    return in_fn

class Timer:
    def __init__(self, timeout:float, cb: Callable[..., Any]) -> None:
        self.timeout:float = timeout
        self.cb = cb

class Task:
    def __init__(self, coro) -> None:
        if not coro:
            return
        self.coro = coro
        f = Future()
        f.set_result(None)
        self.step(f)

    @thread_try
    def step(self, future:Future):
        try:
            netxFuture:Future = self.coro.send(future.result())
        except StopIteration:
            return
        netxFuture.add_done_callback(self.step)

class Thread:
    def __init__(self, worker = 100) -> None:
        self.ready = []
        self.timer = []
        self.timer_map = {}
        self.worker = worker
        self.selector:DefaultSelector = DefaultSelector()
        self.executor = ThreadPoolExecutor(max_workers=worker)

    def __del__(self):
        if self.executor:
            self.executor.shutdown(wait=False)
            self.executor = None

    def fetch(self):
        if len(self.ready):
            return self.ready.pop(0)
        timeout = 0.1
        now = time.time()
        if len(self.timer):
            timeout = self.timer[0].timeout - now
        if timeout < 0:
            timeout = 0.1
        events = self.selector.select(timeout)
        for event_key, event_mask in events:
            self.selector.unregister(event_key.fileobj)
            self.ready.append(event_key.data)
        now = time.time()
        index = 0
        for ti in self.timer:
            #因为排过序的
            if ti.timeout > now:
                break
            index += 1
            self.ready.append(ti.cb)
        if index:
            self.timer = self.timer[index:]

    def run(self):
        while True:
            fn = self.fetch()
            if not fn:
                continue
            try:
                fn()   
            except Exception as e:
                print(f'{fn} err {e}')

    @thread_try
    def add_coroutine(self, coro):
        return Task(coro)

    @thread_try
    def read_await(self, fileobj, cb: Callable[..., Any] = None):
        f = Future()
        def read_wait_cb():
            f.set_result(fileobj)
        self.selector.register(fileobj, EVENT_READ, read_wait_cb)
        ret = yield f
        if cb:
            try:
                cb(ret)
            except Exception as e:
                print(f'{cb} err {e}')
    
    @thread_try
    def socket_remove(self, fileobj):
        self.selector.unregister(fileobj)

    @thread_try
    def write_await(self, fileobj, cb: Callable[..., Any] = None):
        f = Future()
        def write_wait_cb():
            f.set_result(fileobj)
        self.selector.register(fileobj, EVENT_WRITE, write_wait_cb)
        ret = yield f
        if cb:
            try:
                cb(ret)
            except Exception as e:
                print(f'{cb} err {e}')
        return ret

    @thread_try
    def timer_await(self, timeout:float, cb: Callable[..., Any] = None):
        f = Future()
        tkey = (timeout, cb)
        def time_cb():
            del(self.timer_map[tkey])
            f.set_result(timeout)
        tt = Timer(timeout + time.time(), time_cb)
        self.timer_map[tkey] = tt
        self.timer.append(tt)
        #排序
        self.timer.sort(key = lambda x: x.timeout)
        ret = yield f
        if cb:
            try:
                cb(ret)
            except Exception as e:
                print(f'{cb} err {e}')
        return ret
    
    @thread_try
    def timer_remove(self, timeout:float, cb: Callable[..., Any] = None):
        try:
            tkey = (timeout, cb)
            self.timer.remove(self.timer_map[tkey])
            del(self.timer_map[tkey])
        except Exception as e:
                print(f'timer_remove err {e}')

    @thread_try
    def thread_wait(self, fn, cb: Callable[..., Any] = None, *args, **kwargs):
        f = Future()
        def thread_cb(future:Future):
            f.set_result(future.result())
        ef = self.executor.submit(fn, *args, **kwargs)
        ef.add_done_callback(thread_cb)
        ret = yield f
        if cb:
            try:
                cb(ret)
            except Exception as e:
                print(f'{cb} err {e}')
        return ret

_thread = Thread()

def add_coroutine(coro):
    return _thread.add_coroutine(coro)

def thread_run():
    return _thread.run()

def read_await(fileobj, cb: Callable[..., Any] = None):
    return _thread.read_await(fileobj, cb)

def write_await(fileobj, cb: Callable[..., Any] = None):
    return _thread.write_await(fileobj, cb)

def socket_remove(fileobj):
    return _thread.socket_remove(fileobj)

def timer_await(timeout:float, cb: Callable[..., Any] = None):
    return _thread.timer_await(timeout, cb)

def timer_remove(timeout:float, cb: Callable[..., Any] = None):
    return _thread.timer_remove(timeout, cb)

def thread_await(fn, cb: Callable[..., Any] = None, *args, **kwargs):
    return _thread.thread_wait(fn, cb, *args, **kwargs)


def tcp_socket():
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.setblocking(False)
    return sock

def udp_socket():
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.setblocking(False)
    return sock

def tcp_connect(sock:socket.socket, addr:(...)):
    try:
        sock.connect(addr)
    except BlockingIOError:
        return True
    except Exception as e:
        print(f'tcp_connect {e}')
        return False
    return True

def tcp_write(sock:socket.socket, data:bytes):
    try:
        index = 0
        while index < len(data) :
            yield from write_await(sock)
            index += sock.send(data[index:])
        return True
    except Exception as e:
        print(f'tcp_write {e}')
        return False


def tcp_read(sock:socket.socket):
    try:
        response = []
        while True:
            yield from read_await(sock)
            chunk = sock.recv(4096)
            if not chunk:
                break
            response.append(chunk)
        return response
    except Exception as e:
        print(f'tcp_read {e}')
        return []
